package org.ala.tree.rbt;

import java.util.Comparator;
import java.util.LinkedList;

import org.ala.tree.AbstractTree;

/**
 * 红黑树
 	性质1. 结点是红色或黑色。
	性质2. 根结点是黑色。
	性质3. 所有叶子都是黑色。（叶子是NULL结点）（这个性质我也不知道有什么用，最好配上上面的图看，不然会晕）
	性质4. 每个红色结点的两个子结点都是黑色。（从每个叶子到根的所有路径上不能有两个连续的红色结点）
	性质5. 从任一节结点到每个叶子的所有路径都包含相同数目的黑色结点
 *
 * @author ala
 * @date 2025年5月17日
 */
@SuppressWarnings({ "rawtypes", "unchecked" })
public class RBTLinked<K, V> extends AbstractTree<K, V> {
	
	
	/**
	 * 根节点
	 */
	protected Node HEAD = null;
	

	public RBTLinked(Comparator<K> comparator) {
		super(comparator);
	}

	
	@Override
	public boolean add(K k, V v) {
		if (size == 0) {
			HEAD = new Node(k, v);
			//	根节点只能是黑色
			HEAD.red = false;
			size++;
			return true;
		}
		if (_add(HEAD, new Node(k,v))) {size++;}
		HEAD.red = false;
		return true;
	}
	protected boolean _add(Node<K,V> r, Node<K,V> n) {
		int c = comparator.compare(n.key, r.key);
		if (c == 0) {return false;}
		else if (c < 0) {
			if (r.left == NIL) {
				r.left = n; n.parent = r; 
				adjust(n); 
				return true;
			} else {return _add(r.left, n);}
		} else {
			if (r.right == NIL) {
				r.right = n; n.parent = r; 
				adjust(n); 
				return true;
			} else {return _add(r.right, n);}
		}
	}
	/**
	 * 	调整红黑树
	 * 	1 如果插入的是根节点，无需处理（前面的size==0已经判断掉了）
	 *  2 如果父节点是黑，无需处理（新节点是红，不违反规则）
	 *  3 如果父节点是红，需要调整：
	 *  	定义：n（新插入节点）p（父节点）g（爷爷节点）u（叔叔节点）
	 *  	u不存在：
	 *  		n与p同向（n是p的左，p是g的左）
	 *  			旋转g
	 *  			p染黑，g染红（p是旋转后的新根，g和n是其两个子节点）
	 *  		n与p反向（n是p的左，p是g的右。或者对称）
	 *  			反旋p
	 *  			旋转g
	 *  			n染黑（旋转后n是新根，p和g是其两个子节点）
	 *  	u为红：
	 *  		p,u变黑，g变黑
	 *  		g没服务，g是根，g染黑
	 *  		g有父亲，g染红，从g开始往上继续调整
	 *  	u为黑（可以和u不存在的情况合并）：
	 *  		n与p同向
	 *  			旋转g
	 *  			p染黑，g染红
	 *  		n与p反向
	 *  			反旋p
	 *  			旋转g
	 *  			n染黑（旋转后n是新根，p和g是其两个子节点）
	 */
	protected void adjust(Node<K,V> n) {
		Node<K,V> p = n.parent;
		//	如果是根节点或父节点是黑点，不违法规则无需调整
		if (p == null || !p.red) {return;}
		
		//	爷爷和叔叔（父节点是红点一定能拿到爷爷节点，因为根节点永远是黑）
		Node<K,V> g = p.parent,
					   u = g.left == p ? g.right : g.left;
		//	u不存在或u为黑
		if (u == NIL || !u.red) {
			//	n与p同向
			if ((n == p.left && p == g.left) || (n == p.right && p == g.right)) {
				if (n == p.left && p == g.left) { rotationR(g); } 
				else { rotationL(g); }
				p.red = false;
				g.red = true;
			}
			//	n与g反向
			else {
				if (n == p.right && p == g.left) { 
					rotationR(p);
					rotationL(g);
				} else {
					rotationL(p);
					rotationR(g);
				}
				n.red = false;
			}
		}
		//	u为红
		else {
			p.red = u.red = false;
			g.red = false;
			if (g.parent != null) {
				g.red = true;
				adjust(g);
			}
		}
	}
	/**
	 * 	右旋转
	 * 	1）n.right = r
	 *  2）r.left = n.right
	 *  3）n变成新的r
	 */
	protected void rotationR(Node r) {
		Node p = r.parent, n = r.left;
		//	1）n.right = r  2）r.left = n.right
		Node nr = n.right;
		n.right = r;
		r.left = nr; r.parent = n;
		//	3）n变成新的r
		if (p != null) {
			if (r == p.left) {p.left = n;}
			else {p.right = n;}
			n.parent = p;
		} else {
			HEAD = n;
			n.parent = null;
		}
	}
	/**
	 * 	左旋转
	 * 	1）n.left = r
	 *  2）r.right = n.left
	 *  3）n变成新的r
	 */
	protected void rotationL(Node r) {
		Node p = r.parent, n = r.right;
		//	1）n.left = r  2）r.right = n.left
		Node nl = n.left;
		n.left = r;
		r.right = nl; r.parent = n;
		//	3）n变成新的r
		if (p != null) {
			if (r == p.left) {p.left = n;}
			else {p.right = n;}
			n.parent = p;
		} else {
			HEAD = n;
			n.parent = null;
		}
	}


	@Override
	public boolean remove(K k) {
		if (size == 0) {return false;}
		if (HEAD == null) {return false;}
		Node<K,V> n = _search(HEAD, k);
		if (n == null) {return false;}
		if (size == 1) {
			HEAD = null;
			size = 0;
			return true;
		}
		_remove(n);
		size--;
		HEAD.red = false;
		return true;
	}
	/**
	 *	分情况讨论：
	 *	1 删除的是非叶节点
	 *		1.1 删除的节点有左右两个子节点
	 *			找到中序遍历的后继，用后继赋值给当前节点，删除后继（此时后继一定页节点，或只有右子节点）	
	 *		1.2 删除的节点有左或右1个子节点（只可能父是黑，子是红）
	 *			让子节点直接替换自己
	 *			子节点染黑
	 *	2 删除的是叶节点
	 *		2.1 删的是红点
	 *			直接删，不会改变红黑树平衡
	 *		2.2 删的是黑点
	 *			2.2.1 兄弟是黑点：
	 *					2.2.1.1 兄弟的至少有一个红孩子
	 *						r（兄弟的红孩子）s（兄弟）p（父）
	 *						兄弟的红孩子和兄弟同向（r是s的左并且s是p的左，或者对称）
	 *							旋转p
	 *							r染色成s，s染色成p，p染黑
	 *							删除n
	 *						兄弟的红孩子和兄弟反向（r是s的左并且s是p的右，或者对称）
	 *							反旋s
	 *							旋转p
	 *							r染色成p，p染黑
	 *							删除n
	 *					2.2.1.2 兄弟有2个黑儿子
	 *						s（兄弟）p（父）
	 *						兄弟染红
	 *						递归判断p
	 *			2.2.2 兄弟是红点：
	 *				兄父换色
	 *				旋转p
	 *						
	 */
	protected void _remove(Node<K,V> n) {
		//	1.1 删除的节点有左右两个子节点
		if (n.left != NIL && n.right != NIL) {
			Node<K,V> nxt = inorderTraversal(n.key);
			n.key = nxt.key; n.val = nxt.val;
			_remove(nxt);
			return;
		}
		//	1.2 删除的节点有左或右1个子节点
		//	只可能父是黑，子是红。否则违反性质5
		else if (n.left != NIL || n.right != NIL) {
			Node<K,V> p = n.parent, s = n.left != NIL ? n.left : n.right;
			if (p == null) {
				HEAD = s;
			} else {
				if (p.left == n) {p.left = s; s.parent = p;}
				else {p.right = s; s.parent = p;}
				s.red = false;
			}
		}
		//	2 删除的是叶节点
		else {
			//	2.1 删的是红点。直接删，不会改变红黑树平衡
			if (n.red) {
				Node<K,V> p = n.parent;
				if (p.left == n) {p.left = NIL;}
				else {p.right = NIL;}
			}
			//	2.2 删的是黑点
			else {
				Node<K,V> p = n.parent, s = p.left == n ? p.right : p.left;
				//	删除节点
				boolean left = false;
				if (p.left == n) {p.left = NIL; left = true;}
				else {p.right = NIL;}
				
				//	2.2.1 兄弟是黑点：
				if (!s.red) {
					//	调整节点
					adjustRemove(p, left);
				}
				//	2.2.2 兄弟是红点：
				else {
					boolean t = s.red; s.red = p.red; p.red = t;
					if (left) { rotationL(p); }
					else { rotationR(p); }
				}
			}
		}
	}
	/**
	 * 	调整红黑平衡
	 * 	1. 兄弟的红孩子和兄弟同向（r是s的左并且s是p的左，或者对称）
	 * 		旋转p
	 * 		r染色成s，s染色成p，p染黑
	 *  2. 兄弟的红孩子和兄弟反向（r是s的左并且s是p的右，或者对称）
	 *  	反旋s
	 *  	旋转p
	 *  	r染色成p，p染色成黑
	 *  3. 兄弟有2个黑子节点
	 *  	s染成红色
	 *  	p染成黑色
	 *  	如果p不是根节点，从p的父节点开始递归调整
	 * 	
	 * 	@param p 删除节点的父节点
	 *  @param left 删除的是否是左节点
	 */
	protected void adjustRemove(Node<K,V> p, boolean left) {
		if (p == null) {
			return;
		}
		Node<K,V> s = left ? p.right : p.left;
		//	兄弟的红孩子和兄弟同向（r是s的左并且s是p的左，或者对称）
		if ((s.left.red && s == p.left) || (s.right.red && s == p.right)) {
			Node<K,V> r = null;
			if (s.left.red && s == p.left) { r = s.left; rotationR(p); }
			else { r = s.right; rotationL(p); }
			r.red = s.red; s.red = p.red; p.red = false;
			return;
		}
		//	兄弟的红孩子和兄弟反向（r是s的左并且s是p的右，或者对称）
		if ((s.left.red && s == p.right) || (s.right.red && s == p.left)) {
			Node<K,V> r = null;
			if (s.right.red && s == p.left) { r = s.right; rotationL(s); rotationR(p); }
			else { r = s.left; rotationR(s); rotationL(p); }
			r.red = p.red; p.red = false;
			return;
		}
		//	兄弟有2个黑色节点
		if (!s.left.red && !s.right.red) {
			s.red = true;
			p.red = false;
			if (p.parent != null) {
				adjustRemove(p.parent, p.parent.left == p);
			}
		}
	}
	/**
	 * 	中序遍历的下个节点
	 */
	protected Node<K,V> inorderTraversal(K k) {
		Node<K,V> crt = HEAD;
		LinkedList<Node> stack = new LinkedList<>();
		boolean find = false;
		while (crt != null || !stack.isEmpty()) {
			while (crt != null) {
				stack.addLast(crt);
				crt = crt.left;
			}
			crt = stack.poll();
			if (find) {return crt;}
			if (comparator.compare(k, crt.key) == 0) {find = true;}
			crt = crt.right;
		}
		return null;
	}
	

	@Override
	public boolean update(K k, V v) {
		Node<K,V> n = _search(HEAD, k);
		if (n == null) {return false;}
		n.val = v;
		return true;
	}

	@Override
	public V search(K k) {
		Node<K,V> n = _search(HEAD, k);
		return n != null ? n.val : null;
	}
	protected Node<K,V> _search(Node<K,V> r, K k) {
		if (r == NIL) {return null;}
		int c = comparator.compare(k, r.key);
		if (c == 0) {return r;}
		else if (c < 0) {return _search(r.left, k);}
		else {return _search(r.right, k);}
	}

	
	@Override
	public String toString() {
		StringBuilder sbuf = new StringBuilder();
		toString(HEAD, 0, sbuf);
		return sbuf.toString();
	}
	protected void toString(Node<K,V> n, int level, StringBuilder sbuf) {
		if (n == null || n == NIL) {return;}
		StringBuilder t = new StringBuilder();
		for (int i = 0 ; i < level ; i++) {t.append("\t");}
		String lr = "H";
		if (n.parent != null) {
			lr = n.parent.left == n ? "L" : "R";
		}
		sbuf.append(t).append(String.format("(%s,%s)", n.key, n.red ? "R" : "B")).append(lr).append("\n");
		toString(n.left, level+1, sbuf);
		toString(n.right, level+1, sbuf);
	}


	static Node NIL = null;
	static {
		NIL = new Node(null, null);
		NIL.red = false;
		NIL.left = NIL.right = null;
	}
	/**
	 * 	节点信息
	 */
	public static class Node<K, V> {
		K key;
		V val;
		Node parent, left, right;
		boolean red;
		public Node(K key, V val) {
			super();
			this.key = key;
			this.val = val;
			this.red = true;
			this.left = this.right = NIL;
		}
		public String toString() {
			return String.format("(%s,%s)", key, red ? "R" : "H");
		}
	}

}
